#include "ipu_model_wrapper.hpp"
#include <iostream>
#include <stdexcept>
#include <cstring>
#include "mi_ipu.h"

namespace protonet {

// 静态成员变量定义
bool IPUModelWrapper::is_ipu_initialized_ = false;

IPUModelWrapper::IPUModelWrapper(const std::string& model_path, 
                                const std::pair<int, int>& input_size)
    : model_path_(model_path)
    , input_size_(input_size)
    , is_loaded_(false)
    , feature_dimension_(0)
    , channel_id_(0) {
    
    std::cout << "🚀 正在初始化IPU模型包装器..." << std::endl;
    std::cout << "📁 模型路径: " << model_path_ << std::endl;
    std::cout << "📐 输入尺寸: " << input_size_.first << "x" << input_size_.second << std::endl;
    
    if (initializeModel()) {
        is_loaded_ = true;
        std::cout << "✅ IPU模型包装器初始化成功" << std::endl;
    } else {
        std::cerr << "❌ IPU模型包装器初始化失败" << std::endl;
    }
}

IPUModelWrapper::~IPUModelWrapper() {
    if (is_loaded_ && channel_id_ != 0) {
        // 销毁IPU通道
        MI_IPU_DestroyCHN(channel_id_);
        // IPU通道已销毁
    }
    
    // 注意：这里不销毁IPU设备，因为它是全局的静态资源
    // 实际应用中可能需要在程序结束时调用 MI_IPU_DestroyDevice()
}

bool IPUModelWrapper::initializeIPUDevice() {
    if (is_ipu_initialized_) {
        return true; // 已经初始化过了
    }
    
    
    // 初始化MI系统
    MI_S32 ret = MI_SYS_Init(0);
    if (ret != MI_SUCCESS) {
        std::cerr << "❌ MI系统初始化失败，错误码: " << ret << std::endl;
        return false;
    }
    
    // 获取模型的静态信息来确定所需的缓冲区大小
    MI_IPU_OfflineModelStaticInfo_t info;
    ret = MI_IPU_GetOfflineModeStaticInfo(nullptr, const_cast<char*>(model_path_.c_str()), &info);
    if (ret != MI_SUCCESS) {
        std::cerr << "❌ 获取模型静态信息失败，错误码: " << ret << std::endl;
        MI_SYS_Exit(0);
        return false;
    }
    
    // 创建IPU设备
    MI_IPU_DevAttr_t stDevAttr;
    memset(&stDevAttr, 0, sizeof(stDevAttr));
    stDevAttr.u32MaxVariableBufSize = info.u32VariableBufferSize; // 使用模型实际需要的缓冲区大小
    
    ret = MI_IPU_CreateDevice(&stDevAttr, nullptr, nullptr, 0);
    if (ret != MI_SUCCESS) {
        std::cerr << "❌ 创建IPU设备失败，错误码: " << ret << std::endl;
        MI_SYS_Exit(0);
        return false;
    }
    
    is_ipu_initialized_ = true;
    return true;
}

bool IPUModelWrapper::initializeModel() {
    try {
        // 初始化IPU设备
        if (!initializeIPUDevice()) {
            return false;
        }
        
        std::cout << "📂 正在加载模型: " << model_path_ << std::endl;
        
        // 创建IPU通道
        MI_IPUChnAttr_t chn_attr;
        memset(&chn_attr, 0, sizeof(chn_attr));
        
        // 设置通道属性
        chn_attr.u32InputBufDepth = 1;   // 输入缓冲深度
        chn_attr.u32OutputBufDepth = 1;  // 输出缓冲深度
        
        // 创建IPU通道
        MI_S32 ret = MI_IPU_CreateCHN(&channel_id_, &chn_attr, nullptr, 
                                      const_cast<char*>(model_path_.c_str()));
        if (ret != MI_SUCCESS) {
            std::cerr << "❌ 创建IPU通道失败，错误码: " << ret << std::endl;
            return false;
        }
        
        std::cout << "📊 IPU通道创建成功，通道ID: " << channel_id_ << std::endl;
        
        // 获取张量信息
        if (!getTensorInfo()) {
            return false;
        }
        
        return true;
        
    } catch (const std::exception& e) {
        std::cerr << "❌ 模型初始化异常: " << e.what() << std::endl;
        return false;
    }
}

bool IPUModelWrapper::getTensorInfo() {
    try {
        // 获取输入张量信息
        MI_S32 ret = MI_IPU_GetInputTensors(channel_id_, &input_tensor_);
        if (ret != MI_SUCCESS) {
            std::cerr << "❌ 获取输入张量信息失败，错误码: " << ret << std::endl;
            return false;
        }
        
        // 获取输出张量信息
        ret = MI_IPU_GetOutputTensors(channel_id_, &output_tensor_);
        if (ret != MI_SUCCESS) {
            std::cerr << "❌ 获取输出张量信息失败，错误码: " << ret << std::endl;
            return false;
        }
        
        std::cout << "📊 张量信息:" << std::endl;
        std::cout << "   输入张量数量: " << input_tensor_.u32TensorCount << std::endl;
        std::cout << "   输出张量数量: " << output_tensor_.u32TensorCount << std::endl;
        
        if (output_tensor_.u32TensorCount > 0) {
            // 获取输出张量描述来计算特征维度
            MI_IPU_SubNet_InputOutputDesc_t desc;
            ret = MI_IPU_GetInOutTensorDesc(channel_id_, &desc);
            if (ret == MI_SUCCESS && desc.u32OutputTensorCount > 0) {
                // 计算输出特征维度
                feature_dimension_ = 1;
                for (MI_U32 i = 0; i < desc.astMI_OutputTensorDescs[0].u32TensorDim; ++i) {
                    if (desc.astMI_OutputTensorDescs[0].u32TensorShape[i] > 0) {
                        feature_dimension_ *= desc.astMI_OutputTensorDescs[0].u32TensorShape[i];
                    }
                }
                std::cout << "   特征向量维度: " << feature_dimension_ << std::endl;
            } else {
                // 如果无法获取描述，使用默认值
                feature_dimension_ = 512; // 默认特征维度
                std::cout << "   使用默认特征向量维度: " << feature_dimension_ << std::endl;
            }
        }
        
        return true;
        
    } catch (const std::exception& e) {
        std::cerr << "❌ 获取张量信息异常: " << e.what() << std::endl;
        return false;
    }
}

std::vector<float> IPUModelWrapper::infer(const std::vector<float>& input_data) {
    if (!is_loaded_) {
        throw std::runtime_error("模型未正确加载");
    }
    
    try {
        // 获取输入张量描述来计算实际的缓冲区大小
        MI_IPU_SubNet_InputOutputDesc_t desc;
        MI_S32 ret = MI_IPU_GetInOutTensorDesc(channel_id_, &desc);
        if (ret != MI_SUCCESS) {
            throw std::runtime_error("无法获取张量描述，错误码: " + std::to_string(ret));
        }
        
        // 计算实际的输入张量大小（以元素为单位）
        size_t actual_input_elements = 1;
        if (desc.u32InputTensorCount > 0) {
            for (MI_U32 i = 0; i < desc.astMI_InputTensorDescs[0].u32TensorDim; ++i) {
                actual_input_elements *= desc.astMI_InputTensorDescs[0].u32TensorShape[i];
            }
        }
        
        
        // 检查输入数据大小
        if (input_data.size() != actual_input_elements) {
            throw std::runtime_error("输入数据大小不匹配。期望: " + 
                                   std::to_string(actual_input_elements) + 
                                   ", 实际: " + std::to_string(input_data.size()));
        }
        
        // 设置输入数据
        if (input_tensor_.u32TensorCount > 0) {
            MI_IPU_Tensor_t& input_tensor = input_tensor_.astArrayTensors[0];
            
            // 检查张量数据指针
            if (input_tensor.ptTensorData[0] == nullptr) {
                throw std::runtime_error("输入张量缓冲区未分配");
            }
            
            // 复制输入数据到张量
            // IPU张量缓冲区按照元素数量分配，每个元素1字节
            size_t buffer_size = actual_input_elements; // 张量缓冲区大小（字节）
            
            // 将float数据转换为uint8数据复制到张量缓冲区
            // 注意：这里假设模型期望的是量化后的uint8数据
            std::vector<uint8_t> quantized_data(actual_input_elements);
            for (size_t i = 0; i < input_data.size() && i < actual_input_elements; ++i) {
                // 将float值量化到0-255范围（这是一种简单的量化方法）
                float val = input_data[i];
                val = std::max(0.0f, std::min(255.0f, val * 255.0f)); // 假设float在[0,1]范围
                quantized_data[i] = static_cast<uint8_t>(val);
            }
            
            // 尝试复制量化数据到张量缓冲区
            try {
                memcpy(input_tensor.ptTensorData[0], quantized_data.data(), buffer_size);
                MI_SYS_FlushInvCache(input_tensor.ptTensorData[0], buffer_size);
            } catch (...) {
                throw std::runtime_error("内存复制或缓存同步失败");
            }
        } else {
            throw std::runtime_error("输入张量数量为0");
        }
        
        // 执行推理
        ret = MI_IPU_Invoke(channel_id_, &input_tensor_, &output_tensor_);
        if (ret != MI_SUCCESS) {
            throw std::runtime_error("IPU推理执行失败，错误码: " + std::to_string(ret));
        }
        
        // 获取输出数据
        std::vector<float> features;
        if (output_tensor_.u32TensorCount > 0) {
            MI_IPU_Tensor_t& output_tensor = output_tensor_.astArrayTensors[0];
            
            // 检查输出张量数据指针
            if (output_tensor.ptTensorData[0] == nullptr) {
                throw std::runtime_error("输出张量缓冲区未分配");
            }
            
            // 复制输出数据
            float* output_data = reinterpret_cast<float*>(output_tensor.ptTensorData[0]);
            size_t output_size = feature_dimension_;
            
            features.assign(output_data, output_data + output_size);
        } else {
            throw std::runtime_error("输出张量数量为0");
        }
        
        return features;
        
    } catch (const std::exception& e) {
        throw std::runtime_error("推理过程中发生错误: " + std::string(e.what()));
    }
}

} // namespace protonet